/**
 * @file   MultigridSolver.h
 * @author WinterMelonCDY <wintermeloncdy@wintermeloncdy-VirtualBox>
 * @date   Fri Apr 21 14:47:54 2023
 * 
 * @brief  
 * 
 * 
 */

#include <iostream>
#include <cmath>
#include <algorithm>
#include <vector>
#include <eigen3/Eigen/Dense>
#include <eigen3/Eigen/Sparse>
#include <eigen3/Eigen/Core>

template <int dim>
class Function
{
public:
    virtual double operator()(double _x) = 0;       // 返回函数u(_x)
    virtual double diff(double _x) { return 0; }    // 返回函数u(_x)导数
    virtual double Laplace(double _x) { return 0; } // 返回函数u(_x)经过-Laplace算子作用后的值（二阶导取负）
};

template <int dim>
class Weighted_Jacobi //加权Jacobi迭代
{
private:
    Eigen::SparseMatrix<double> A; // 迭代方程为Ax=b
    Eigen::VectorXd b;
    Eigen::VectorXd init_x; // 迭代初始值
    int iter_Num;           // 迭代次数
    double w;               // 迭代权重
public:
    Weighted_Jacobi(Eigen::SparseMatrix<double> _A, Eigen::VectorXd _b, Eigen::VectorXd _init_x, int _iter_Num, double _w) : A(_A), b(_b), init_x(_init_x), iter_Num(_iter_Num), w(_w) {}

    // 迭代函数
    Eigen::VectorXd solve()
    {

        Eigen::SparseMatrix<double> D(A.rows(), A.cols());
        Eigen::SparseMatrix<double> D_inverse(A.rows(), A.cols());

        for (int i = 0; i < A.rows(); i++)
        {
            D.insert(i, i) = A.coeff(i, i);
            D_inverse.insert(i, i) = 1.0 / A.coeff(i, i);
        }

        Eigen::SparseMatrix<double> T = D_inverse * (D - A);
        Eigen::VectorXd c = D_inverse * b;
        Eigen::VectorXd x = init_x;

        for (int k = 0; k < iter_Num; k++)
        {
            Eigen::VectorXd x_star = T * x + c;
            x = (1.0 - w) * x + w * x_star;
        }

        return x;
    }
};

template <int dim>
class Transformation
{
private:
    Eigen::VectorXd v; // 待处理向量
    double h;          
public:
    Transformation(Eigen::VectorXd _v, double _h) : v(_v), h(_h) {}

    // Interpolation operator，分为linear和quadratic
    Eigen::VectorXd Interpolation(std::string inter, std::string boundary)
      {
	if (inter == "Linear")
        {
            int n = static_cast<int>(2.0 / h + 1.0);
            Eigen::SparseMatrix<double> I(n, (n + 1) / 2);

            for (int i = 1; i < (n - 1) / 2; i++)
            {
                I.insert(2 * i, i) = 1.0;
                I.insert(2 * i - 1, i) = 0.5;
		I.insert(2 * i + 1, i) = 0.5;
            }
                I.insert(0, 0) = 1.0;
                I.insert(1, 0) = 0.5;
                I.insert(n-1, (n-1)/2) = 1.0;
                I.insert(n-2, (n-1)/2) = 0.5;

		return I * v;
	}
        if (inter == "Quadratic")
        {
            int n = static_cast<int>(2.0 / h + 1.0);
            Eigen::SparseMatrix<double> I(n, (n + 1) / 2);

            for (int i = 0; i < (n + 1) / 2; i++)
            {
                I.insert(2 * i, i) = 1.0;
            }
            for (int i = 1; i < (n - 3) / 2; i++)
            {
                I.insert(2 * i + 1, i) = 9.0/16.0;
                I.insert(2 * i + 1, i - 1) = -1.0/16.0;
                I.insert(2 * i + 1, i + 1) = 9.0/16.0;
                I.insert(2 * i + 1, i + 2) = -1.0/16.0;
            }
                I.insert(1, 1) = 0.5;
                I.insert(1, 0) = 0.5;
                I.insert(n-2, (n-3)/2) = 0.5;
                I.insert(n-2, (n-1)/2) = 0.5;

                return I * v;
	}
        Eigen::VectorXd vec;
        return vec;
      }

    // Restriction operator，分为full weighting和injection
    Eigen::VectorXd Restriction(std::string res, std::string boundary)
      {
	if (res == "FO")
	  {
	    int n = static_cast<int>(1.0 / (2.0 * h) + 1.0);
	    Eigen::SparseMatrix<double> I(n, n * 2 - 1);
	    for (int i = 1; i < n - 1; i++)
	      {
		I.insert(i, 2 * i) = 0.5;
                I.insert(i, 2 * i - 1) = 0.25;
                I.insert(i, 2 * i + 1) = 0.25;
	      }
                I.insert(0, 0) = 1.0;
                I.insert(n - 1, n * 2 - 2) = 1.0;           

                return I * v;
	  }
	if (res == "IO")
	  {
	    int n = static_cast<int>(1.0 / (2.0 * h) + 1.0);
	    Eigen::SparseMatrix<double> I(n, n * 2 - 1);
	    for (int i = 0; i < n ; i++)
              {
                I.insert(i, 2 * i) = 1.0;
              }

                return I * v;
	  }
	Eigen::VectorXd vec;
	return vec;
      }
};

template <int dim>
class Coefficient
{
private:
    Function<dim> &u;     // 给定函数
    double h;             // 网格宽度
    std::string boundary; // 边界条件
public:
    Coefficient(Function<dim> &_u, std::string _boundary, double _h) : u(_u), boundary(_boundary), h(_h) {}

    // 构造矩阵A_h，分三种边界条件进行讨论
    Eigen::SparseMatrix<double> get_Ah()
    {
      if (boundary == "DC")
	{
	  int n = static_cast<int>(1.0 / h + 1.0);
          Eigen::SparseMatrix<double> Ah(n, n);
          std::vector<Eigen::Triplet<double> > tripletlist;

          for (int i = 1; i < n - 1; ++i)
          {
            tripletlist.push_back(Eigen::Triplet<double>(i, i, 2.0/(h*h)));
            tripletlist.push_back(Eigen::Triplet<double>(i, i-1, -1.0/(h*h)));
            tripletlist.push_back(Eigen::Triplet<double>(i, i+1, -1.0/(h*h)));
          }
            tripletlist.push_back(Eigen::Triplet<double>(0, 0, 1.0));
            tripletlist.push_back(Eigen::Triplet<double>(n-1, n-1, 1.0));
            Ah.setFromTriplets(tripletlist.begin(), tripletlist.end());
            Ah.makeCompressed();

            return Ah;
	}
      if (boundary == "NC")
        {
          int n = static_cast<int>(1.0 / h + 1.0);
          Eigen::SparseMatrix<double> Ah(n, n);
          std::vector<Eigen::Triplet<double> > tripletlist;

          for (int i = 1; i < n - 1; ++i)
          {
	    tripletlist.push_back(Eigen::Triplet<double>(i, i, 2.0/(h*h)));
            tripletlist.push_back(Eigen::Triplet<double>(i, i-1, -1.0/(h*h)));
            tripletlist.push_back(Eigen::Triplet<double>(i, i+1, -1.0/(h*h)));
          }
            tripletlist.push_back(Eigen::Triplet<double>(0, 0, 1.0/h));
            tripletlist.push_back(Eigen::Triplet<double>(0, 1, 2.0/(h)));
            tripletlist.push_back(Eigen::Triplet<double>(0, 2, -1.0/(2.0*h)));
            tripletlist.push_back(Eigen::Triplet<double>(n-1, n-1, 3.0/(2.0*h)));
            tripletlist.push_back(Eigen::Triplet<double>(n-1, n-2, -2.0/(h)));
            tripletlist.push_back(Eigen::Triplet<double>(n-1, n-3, 1.0/(2.0*h)));
            Ah.setFromTriplets(tripletlist.begin(), tripletlist.end());
            Ah.makeCompressed();

            return Ah;
	}
      if (boundary == "MC")
	{
	  int n = static_cast<int>(1.0 / h + 1.0);
          Eigen::SparseMatrix<double> Ah(n, n);
          std::vector<Eigen::Triplet<double> > tripletlist;

          for (int i = 1; i < n - 1; ++i)
	  {
            tripletlist.push_back(Eigen::Triplet<double>(i, i, 2.0/(h*h)));
            tripletlist.push_back(Eigen::Triplet<double>(i, i-1, -1.0/(h*h)));
            tripletlist.push_back(Eigen::Triplet<double>(i, i+1, -1.0/(h*h)));
          }
            tripletlist.push_back(Eigen::Triplet<double>(0, 0, 1.0));
            tripletlist.push_back(Eigen::Triplet<double>(n-1, n-1, 3.0/(2.0*h)));
            tripletlist.push_back(Eigen::Triplet<double>(n-1, n-2, -2.0/(h)));
            tripletlist.push_back(Eigen::Triplet<double>(n-1, n-3, 1.0/(2.0*h)));
            Ah.setFromTriplets(tripletlist.begin(), tripletlist.end());
            Ah.makeCompressed();

            return Ah;
	}
        Eigen::SparseMatrix<double> vec;
        return vec;
    }

    // 构造初始右端向量，分三种边界条件进行讨论
    Eigen::VectorXd get_initf()
    {
      if (boundary == "DC")
	{
          int n = static_cast<int>(1.0 / h + 1.0);
          Eigen::VectorXd f_init(n);

          f_init(0) = u(0);
          f_init(n - 1) = u(1);

          for (int i = 1; i < n - 1; i++)
	    {
              f_init(i) = u.Laplace(i * h);
            }

	  return f_init;
	}
      if (boundary == "NC")
	{
          int n = static_cast<int>(1.0 / h + 1.0);
          Eigen::VectorXd f_init(n);

          f_init(0) = u.diff(0) + (3.0*u(0))/(2.0*h) + (1.0*u(0))/h;
          f_init(n - 1) = u.diff(1);

          for (int i = 1; i < n - 1; i++)
	    {
              f_init(i) = u.Laplace(i * h);
            }

	  return f_init;
	}
      if (boundary == "MC")
	{
          int n = static_cast<int>(1.0 / h + 1.0);
          Eigen::VectorXd f_init(n);

          f_init(0) = u(0);
          f_init(n - 1) = u.diff(1);

          for (int i = 1; i < n - 1; i++)
            {
              f_init(i) = u.Laplace(i * h);
            }

          return f_init;
	}
      Eigen::VectorXd vec;
      return vec;
    }
};


template <int dim>
class Solver
{
private:
    int v1, v2;
    Function<dim> &u;          // 网格函数
    std::string boundary;      // Boundary condition
    std::string resop;         // Restriction operator
    std::string interop;       // Interpolation operator
    int maxiter;               // Maximum iteration number
    double relacc;             // Relative accuracy
    double h;                  // 网格宽度
    std::vector<double> initv; // 迭代初始向量
    Eigen::VectorXd v;         // 计算解
    std::vector<Eigen::SparseMatrix<double>> A; // 预先构建Ah
    std::vector<double> Re;    // 存储residual

    // Vcycle求解
    Eigen::VectorXd solve_Vcycle(Eigen::VectorXd v_h, Eigen::VectorXd f_h, double H)
    {
        Eigen::SparseMatrix<double> A_h = A[log2(static_cast<int>(H/h))];
        Weighted_Jacobi<dim> J1(A_h, f_h, v_h, v1, 2.0/3.0);
        v_h = J1.solve();
        if (H >= 0.5)
        {
            Weighted_Jacobi<dim> J2(A_h, f_h, v_h, v2, 2.0/3.0);
            return J2.solve();
        }
        else
        {
            Transformation<dim> T1(f_h - A_h * v_h, H);
            Eigen::VectorXd f_2h = T1.Restriction(resop, boundary);
            Eigen::VectorXd v_2h = Eigen::VectorXd::Zero(f_2h.size());
            v_2h = solve_Vcycle(v_2h, f_2h, 2 * H);
            Transformation<dim> T2(v_2h, 2 * H);
            v_h = v_h + T2.Interpolation(interop, boundary);
            Weighted_Jacobi<dim> J2(A_h, f_h, v_h, v2, 2.0/3.0);
            return J2.solve();
        }
    }

    // FMG求解
    Eigen::VectorXd solve_FMG(Eigen::VectorXd f_h, double H)
    {
        Eigen::VectorXd v_h;
        if (H >= 0.5)
        {
            v_h = Eigen::VectorXd::Zero(f_h.size());
            v_h = solve_Vcycle(v_h, f_h, H);
            return v_h;
        }
        else
        {
            Transformation<dim> T1(f_h, H);
            Eigen::VectorXd f_2h = T1.Restriction(resop, boundary);
            Eigen::VectorXd v_2h = Eigen::VectorXd::Zero(f_2h.size());
            v_2h = solve_FMG(f_2h, 2 * H);
            Transformation<dim> T2(v_2h, 2 * H);
            v_h = T2.Interpolation(interop, boundary);
            v_h = solve_Vcycle(v_h, f_h, H);

            return v_h;
        }    
    }


public:
    Solver(Function<dim> &_u,
            std::string _boundary,
            std::string _resop,
            std::string _interop,
            double _maxiter,
            double _relacc,
            double _h,
            std::vector<double> _initv)
        : u(_u),
          boundary(_boundary),
          resop(_resop),
          interop(_interop),
          maxiter(_maxiter),
          relacc(_relacc),
          h(_h),
          initv(_initv)
	  {
	    v1 = 5;
	    v2 = 5;

	    double h_A = h;
	    while (h_A <= 0.5)
	      {
		Coefficient<dim> C(u, boundary, h_A);
                Eigen::SparseMatrix<double> A_h = C.get_Ah();
                A.push_back(A_h);
                h_A = h_A*2;
              }
	  }
    
    void Vcycle()
    {
        Eigen::VectorXd init_f;
        Eigen::SparseMatrix<double> init_A;
        Coefficient<dim> C(u, boundary, h);

        init_f = C.get_initf();
        init_A = C.get_Ah();
        Eigen::VectorXd init_v = Eigen::VectorXd::Map(initv.data(), initv.size());
        v = init_v;

        int iter = 0;
        double rel_res_norm = 1.0;

        while (iter < maxiter && rel_res_norm >= relacc)
        {
            v = solve_Vcycle(v, init_f, h);
            Eigen::VectorXd r = init_f - init_A * v;
            rel_res_norm = r.norm() / init_f.norm();
	    Re.push_back(h*r.lpNorm<1>());
            iter++;
        }

    }

    void FMG()
    {
        Eigen::VectorXd init_f;
        Eigen::VectorXd f;
        Eigen::VectorXd v_temp;
        Eigen::SparseMatrix<double> init_A;
        Coefficient<dim> C(u, boundary, h);

        init_f = C.get_initf();
        init_A = C.get_Ah();
        v = Eigen::VectorXd::Zero(init_f.size());

        int iter = 0;
        double rel_res_norm = 1.0;
        v = solve_Vcycle(v, init_f, h);
        f = init_f - init_A*v;

        while (iter < maxiter && rel_res_norm >= relacc)
        {
            v_temp = solve_FMG(f, h);
            v = v + v_temp;
            f = f - init_A * v_temp;
            rel_res_norm = f.norm() / init_f.norm();
            Re.push_back(sqrt(h)*f.norm());
            iter++;
        }
    }

    std::vector<double> get_residual() // 返回每一次迭代残差的一范数
    {
        return Re;
    }
    double get_errornorm() // 返回误差的无穷范数
    {
        Eigen::VectorXd U_hat = Eigen::VectorXd::Zero(v.size()); // U_hat存储真实值

        for (int i = 0; i < v.size(); ++i)
	  {
	    U_hat(i) = u(i*h);
          }
	double normInf = (v - U_hat).lpNorm<Eigen::Infinity>();	
	return normInf;
    }
};
